import json
import ipaddress
from openpyxl import Workbook
from openpyxl.styles import Font, Alignment
from openpyxl.worksheet.hyperlink import Hyperlink
import argparse
import nmap
import re
import csv
import requests
from io import StringIO
import os

# 下载并加载 IANA 的端口服务分配表
def load_iana_port_mapping():
    url = "https://www.iana.org/assignments/service-names-port-numbers/service-names-port-numbers.csv"
    local_file = "service-names-port-numbers.csv"
    iana_mapping = {}

    try:
        # 尝试从 IANA 网站下载 CSV 文件
        response = requests.get(url)
        response.raise_for_status()
        
        # 将下载的内容保存到本地文件
        with open(local_file, 'w', encoding='utf-8') as file:
            file.write(response.text)
        print(f"Successfully downloaded and saved IANA port mapping to {local_file}.")
        
        # 从本地文件加载数据
        with open(local_file, 'r', encoding='utf-8') as file:
            csv_data = StringIO(file.read())
    except Exception as e:
        print(f"Failed to download IANA port mapping: {e}")
        # 如果下载失败，尝试读取本地文件
        if os.path.exists(local_file):
            try:
                with open(local_file, 'r', encoding='utf-8') as file:
                    csv_data = StringIO(file.read())
                print("Loaded IANA port mapping from local file.")
            except Exception as e:
                print(f"Failed to load local IANA port mapping: {e}")
                return iana_mapping
        else:
            print("Local IANA port mapping file does not exist. Skipping IANA port mapping.")
            return iana_mapping

    # 解析 CSV 文件
    reader = csv.DictReader(csv_data)
    for row in reader:
        service_name = row['Service Name'].lower()
        port_number = row['Port Number']
        if service_name and port_number and port_number.isdigit():
            if service_name not in iana_mapping:
                iana_mapping[service_name] = set()
            iana_mapping[service_name].add(port_number)
    return iana_mapping

# 加载NVD JSON文件
def load_nvd_json(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        return json.load(file)

# 从NVD数据中提取关键字
def extract_keywords_from_nvd(nvd_data):
    keywords = set()
    for cve_item in nvd_data['CVE_Items']:
        # 提取CVE描述中的关键字
        description = cve_item.get('cve', {}).get('description', {}).get('description_data', [])
        for desc in description:
            for word in desc.get('value', '').split():
                if len(word) > 3:  # 过滤掉过短的词
                    keywords.add(word.lower())
        
        # 提取受影响的产品信息
        affects = cve_item.get('cve', {}).get('affects', {})
        vendor_data = affects.get('vendor', {}).get('vendor_data', [])
        for vendor in vendor_data:
            for product in vendor.get('product', {}).get('product_data', []):
                product_name = product.get('product_name', '').lower()
                if product_name:
                    keywords.add(product_name)
        
        # 提取配置信息中的关键字
        configurations = cve_item.get('configurations', {}).get('nodes', [])
        for node in configurations:
            for cpe_match in node.get('cpe_match', []):
                cpe_uri = cpe_match.get('cpe23Uri', '').lower()
                if cpe_uri:
                    keywords.add(cpe_uri)
    return list(keywords)

# 检查CVE是否影响指定的软件或服务
def check_cve_for_software(cve_item, software_keywords):
    # 检查CVE描述中是否包含指定的软件或服务关键字
    description = cve_item.get('cve', {}).get('description', {}).get('description_data', [])
    for desc in description:
        for keyword in software_keywords:
            if keyword.lower() in desc.get('value', '').lower():
                return True

    # 检查受影响的产品信息
    affects = cve_item.get('cve', {}).get('affects', {})
    vendor_data = affects.get('vendor', {}).get('vendor_data', [])
    for vendor in vendor_data:
        for product in vendor.get('product', {}).get('product_data', []):
            for keyword in software_keywords:
                if keyword.lower() in product.get('product_name', '').lower():
                    return True

    # 检查配置信息中的受影响软件
    configurations = cve_item.get('configurations', {}).get('nodes', [])
    for node in configurations:
        for cpe_match in node.get('cpe_match', []):
            for keyword in software_keywords:
                if keyword.lower() in cpe_match.get('cpe23Uri', '').lower():
                    return True

    return False

# 模拟主机扫描，动态生成关键字列表
def get_keywords_for_ip(ip, nvd_keywords):
    # 使用 nmap 扫描 IP 地址的开放端口
    nm = nmap.PortScanner()
    nm.scan(ip, arguments='-sV')  # 扫描 IP 地址的开放端口和服务版本
    
    # 提取扫描结果中的服务信息
    services = set()
    for host in nm.all_hosts():
        for proto in nm[host].all_protocols():
            ports = nm[host][proto].keys()
            for port in ports:
                service = nm[host][proto][port]['name']
                services.add(service)
    
    # 将服务信息与 NVD 关键字进行匹配
    keywords = []
    for service in services:
        for keyword in nvd_keywords:
            if service.lower() in keyword.lower():
                keywords.append(keyword)
    
    return keywords

# 获取风险等级
def get_risk_level(cvss_score):
    if cvss_score >= 9.0:
        return "高危"
    elif cvss_score >= 7.0:
        return "中危"
    elif cvss_score >= 4.0:
        return "低危"
    else:
        return "信息"

# 获取加固建议
def get_remediation_advice(cve_id):
    # 这里可以根据 CVE ID 返回具体的加固建议
    # 例如，从外部数据库或文件中获取建议
    return "请参考官方补丁或升级到最新版本。"

# 获取应用分类
def get_application_category(software_keywords):
    # 更细化的应用分类
    if "openssh" in software_keywords:
        return "SSH服务漏洞"
    elif "apache" in software_keywords:
        return "Apache HTTP Server漏洞"
    elif "nginx" in software_keywords:
        return "NGINX Web服务器漏洞"
    elif "mysql" in software_keywords:
        return "MySQL数据库漏洞"
    elif "oracle" in software_keywords:
        return "Oracle数据库漏洞"
    elif "wordpress" in software_keywords:
        return "WordPress应用漏洞"
    elif "ftp" in software_keywords:
        return "FTP服务漏洞"
    elif "tls" in software_keywords or "ssl" in software_keywords:
        return "TLS/SSL协议漏洞"
    elif "php" in software_keywords:
        return "PHP语言漏洞"
    elif "python" in software_keywords:
        return "Python语言漏洞"
    elif "java" in software_keywords:
        return "Java语言漏洞"
    elif "docker" in software_keywords:
        return "Docker容器漏洞"
    elif "kubernetes" in software_keywords:
        return "Kubernetes集群漏洞"
    elif "linux kernel" in software_keywords:
        return "Linux内核漏洞"
    elif "windows" in software_keywords:
        return "Windows系统漏洞"
    elif "web server" in software_keywords:
        return "Web服务器漏洞"
    elif "web application" in software_keywords:
        return "Web应用漏洞"
    elif "unix" in software_keywords:
        return "Unix系统漏洞"
    elif "linux" in software_keywords:
        return "Linux系统漏洞"
    else:
        return "其他漏洞"

# 获取漏洞名称
def get_vulnerability_name(cve_item):
    cve_id = cve_item.get('cve', {}).get('CVE_data_meta', {}).get('ID', '')
    description = cve_item.get('cve', {}).get('description', {}).get('description_data', [])
    description_text = ' '.join([desc.get('value', '') for desc in description]).lower()
    
    # 根据描述文本动态生成漏洞名称
    if "openssh" in description_text:
        return f"OpenSSH 安全漏洞({cve_id})"
    elif "mysql" in description_text or "mariadb" in description_text:
        return f"MySQL 安全漏洞({cve_id})"
    elif "apache" in description_text or "httpd" in description_text:
        return f"Apache HTTP Server 安全漏洞({cve_id})"
    elif "nginx" in description_text:
        return f"NGINX 安全漏洞({cve_id})"
    elif "wordpress" in description_text:
        return f"WordPress 安全漏洞({cve_id})"
    elif "oracle" in description_text:
        return f"Oracle 安全漏洞({cve_id})"
    elif "linux kernel" in description_text or "linux" in description_text:
        return f"Linux 内核安全漏洞({cve_id})"
    elif "windows" in description_text:
        return f"Windows 安全漏洞({cve_id})"
    elif "ssh" in description_text:
        return f"SSH 安全漏洞({cve_id})"
    elif "tls" in description_text or "ssl" in description_text:
        return f"TLS/SSL 安全漏洞({cve_id})"
    elif "php" in description_text:
        return f"PHP 安全漏洞({cve_id})"
    elif "python" in description_text:
        return f"Python 安全漏洞({cve_id})"
    elif "java" in description_text:
        return f"Java 安全漏洞({cve_id})"
    elif "docker" in description_text:
        return f"Docker 安全漏洞({cve_id})"
    elif "kubernetes" in description_text:
        return f"Kubernetes 安全漏洞({cve_id})"
    elif "ftp" in description_text or "file transfer protocol" in description_text:
        return f"FTP 安全漏洞({cve_id})"
    elif "web server" in description_text or "web application" in description_text:
        return f"Web 应用安全漏洞({cve_id})"
    elif "unix" in description_text or "linux" in description_text:
        return f"Unix/Linux 安全漏洞({cve_id})"
    elif "windows" in description_text:
        return f"Windows 安全漏洞({cve_id})"
    else:
        # 如果未匹配到特定分类，返回通用名称
        return f"安全漏洞({cve_id})"

# 从CVE描述中提取端口号
def extract_ports_from_description(description_text):
    # 使用正则表达式提取端口号
    port_pattern = r'port\s+(\d{1,5})|(\d{1,5})\/tcp|(\d{1,5})\/udp'
    ports = set()
    matches = re.findall(port_pattern, description_text, re.IGNORECASE)
    for match in matches:
        for port in match:
            if port and 0 <= int(port) <= 65535:
                ports.add(port)
    return list(ports)

# 获取影响端口
def get_affected_ports(cve_item, nm, iana_mapping):
    affected_ports = set()
    
    # 从CVE描述中提取端口号
    description = cve_item.get('cve', {}).get('description', {}).get('description_data', [])
    description_text = ' '.join([desc.get('value', '') for desc in description])
    ports_from_description = extract_ports_from_description(description_text)
    affected_ports.update(ports_from_description)
    
    # 从CVE配置信息中提取服务名称
    configurations = cve_item.get('configurations', {}).get('nodes', [])
    for node in configurations:
        for cpe_match in node.get('cpe_match', []):
            cpe23Uri = cpe_match.get('cpe23Uri', '')
            service_name = cpe23Uri.split(':')[4].lower()  # 提取服务名称
            if service_name in iana_mapping:
                affected_ports.update(iana_mapping[service_name])
    
    # 从Nmap扫描结果中匹配服务
    for host in nm.all_hosts():
        for proto in nm[host].all_protocols():
            ports = nm[host][proto].keys()
            for port in ports:
                service = nm[host][proto][port]['name'].lower()
                if service in iana_mapping:
                    affected_ports.update(iana_mapping[service])
    
    return affected_ports

# 扫描漏洞并保存结果
def scan_vulnerabilities(nvd_data, ip_range, output_file, iana_mapping):
    # 创建一个新的Excel工作簿
    wb = Workbook()
    ws = wb.active
    ws.title = "Vulnerability Scan Results"
    
    # 添加表头
    headers = ["序号", "风险等级", "影响IP", "影响端口", "漏洞名称", "CVE ID", "CVE_ID信息", "CNNVD ID", "CNVD ID", "CNCVE ID", "CVSS分值", "检出方式", "出现次数", "应用分类", "漏洞描述", "加固建议"]
    ws.append(headers)
    
    # 设置表头为筛选状态
    ws.auto_filter.ref = ws.dimensions
    
    # 从NVD数据中提取关键字
    nvd_keywords = extract_keywords_from_nvd(nvd_data)
    print(f"Extracted {len(nvd_keywords)} keywords from NVD data.")
    
    # 用于统计漏洞出现次数
    cve_count = {}
    
    # 遍历IP段中的所有IP地址
    for ip in ip_range:
        ip_str = str(ip)
        print(f"\nScanning IP: {ip_str}")  # 打印当前扫描的IP地址
        
        # 动态生成关键字列表（最大匹配）
        software_keywords = get_keywords_for_ip(ip_str, nvd_keywords)
        print(f"Using keywords: {software_keywords}")  # 打印当前使用的关键字
        
        # 使用 nmap 扫描 IP 地址的开放端口
        nm = nmap.PortScanner()
        nm.scan(ip_str, arguments='-sV')  # 扫描 IP 地址的开放端口和服务版本
        
        # 检查每个CVE是否影响当前IP
        for cve_item in nvd_data['CVE_Items']:
            if check_cve_for_software(cve_item, software_keywords):
                cve_id = cve_item.get('cve', {}).get('CVE_data_meta', {}).get('ID', '')
                description = cve_item.get('cve', {}).get('description', {}).get('description_data', [])
                description_text = ' '.join([desc.get('value', '') for desc in description])
                
                # 获取CVSS评分
                cvss_score = float(cve_item.get('impact', {}).get('baseMetricV3', {}).get('cvssV3', {}).get('baseScore', 0.0))
                risk_level = get_risk_level(cvss_score)
                
                # 获取影响端口
                affected_ports = get_affected_ports(cve_item, nm, iana_mapping)
                
                # 统计漏洞出现次数
                if cve_id in cve_count:
                    cve_count[cve_id] += 1
                else:
                    cve_count[cve_id] = 1
                
                # 获取加固建议
                remediation_advice = get_remediation_advice(cve_id)
                
                # 获取应用分类
                application_category = get_application_category(software_keywords)
                
                # 获取漏洞名称
                vulnerability_name = get_vulnerability_name(cve_item)
                
                # 生成 CVE_ID信息 链接
                cve_info_url = f"https://cve.mitre.org/cgi-bin/cvekey.cgi?keyword={cve_id}"
                
                # 将结果添加到Excel中
                row = [
                    len(ws['A']),  # 序号
                    risk_level,  # 风险等级
                    ip_str,  # 影响IP
                    ', '.join(affected_ports),  # 影响端口
                    vulnerability_name,  # 漏洞名称
                    cve_id,  # CVE ID
                    cve_info_url,  # CVE_ID信息
                    "",  # CNNVD ID（需从外部数据获取）
                    "",  # CNVD ID（需从外部数据获取）
                    "",  # CNCVE ID（需从外部数据获取）
                    cvss_score,  # CVSS分值
                    "自动扫描",  # 检出方式
                    cve_count[cve_id],  # 出现次数
                    application_category,  # 应用分类
                    description_text,  # 漏洞描述
                    remediation_advice  # 加固建议
                ]
                ws.append(row)
                
                # 设置 CVE_ID信息 列为超链接
                cve_info_cell = ws.cell(row=ws.max_row, column=headers.index("CVE_ID信息") + 1)
                cve_info_cell.hyperlink = cve_info_url
                cve_info_cell.value = "查看详情"
                cve_info_cell.font = Font(color="0000FF", underline="single")
                cve_info_cell.alignment = Alignment(horizontal="left")
    
    # 保存Excel文件
    wb.save(output_file)
    print(f"\nScan results saved to {output_file}")

# 主函数
def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description="Vulnerability Scanner")
    parser.add_argument("--ip-range", required=True, help="IP range to scan (e.g., 192.168.6.0/24)")
    parser.add_argument("--nvd-file", default="nvdcve-1.1-modified.json", help="Path to NVD JSON file")
    parser.add_argument("--output-file", default="vulnerability_scan_results.xlsx", help="Path to output Excel file")
    args = parser.parse_args()

    # 加载NVD数据
    nvd_data = load_nvd_json(args.nvd_file)
    
    # 解析IP段
    ip_range = ipaddress.ip_network(args.ip_range)
    
    # 加载 IANA 端口服务分配表
    iana_mapping = load_iana_port_mapping()
    
    # 扫描漏洞并保存结果
    scan_vulnerabilities(nvd_data, ip_range, args.output_file, iana_mapping)

if __name__ == "__main__":
    main()
